package com.jiazhong.spring.ai.elasticsearch.test;

import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.junit.jupiter.api.Test;
import org.springframework.ai.document.Document;
import org.springframework.ai.reader.TextReader;
import org.springframework.ai.reader.markdown.MarkdownDocumentReader;
import org.springframework.ai.reader.markdown.config.MarkdownDocumentReaderConfig;
import org.springframework.ai.reader.pdf.PagePdfDocumentReader;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.test.context.SpringBootTest;

import java.util.List;

@SpringBootTest
@Slf4j
public class App {
    @Resource
    private VectorStore vectorStore;
    @Value("classpath:text-source.txt")
    private org.springframework.core.io.Resource txtResource;
    @Value("classpath:markdown-source.md")
    private org.springframework.core.io.Resource markdownResource;
    @Value("classpath:pdf-source.pdf")
    private org.springframework.core.io.Resource pdfResource;

    @Test
    public void add1() {
        List<Document> documents = List.of(
                new Document("I like SpringAI"),
                new Document("I like SpringBoot"),
                new Document("I like ElasticSearch"),
                new Document("I like Spring ElasticSearch"),
                new Document("I love ElasticSearch"),
                new Document("I like Python")

        );
        vectorStore.add(documents);
    }

    @Test
    public void remove() {
        List<String> ids = List.of("a9dc75eb-9643-4310-bd09-42c47e8e03e8", "d92e70bf-ec7c-4a6a-b12e-b56e6b982d4d",
                "4a4cf6ff-92c4-4ec3-9502-9c8cdf8a914e", "0eb5482f-0183-4f3d-9709-549a3e3e1c50",
                "add4855d-2d81-47af-a9af-510c3c341c00", "9d2a9dff-18bf-4917-8dc8-1ab8eb8f143a");
        vectorStore.delete(ids);
    }

    @Test
    public void find1() {
        List<Document> documents = vectorStore.similaritySearch("九月西安穿什么");
        documents.forEach(System.out::println);
    }

    @Test
    public void fin2() {
        SearchRequest searchRequest = SearchRequest.builder()
                .query("I like")
                .topK(6)
                .similarityThreshold(0.5)
                .build();
        List<Document> documents = vectorStore.similaritySearch(searchRequest);
        documents.forEach(System.out::println);
    }

    @Test
    public void add2() {
        TextReader txtReader = new TextReader(txtResource);
        List<Document> documents = txtReader.get();
        // 转换：按Token拆分
        TokenTextSplitter splitter = new TokenTextSplitter();
        List<Document> chunks = splitter.apply(documents);
        // 加载: 存储到ES向量数据库
        vectorStore.add(chunks);
    }

    @Test
    public void add3() {
        PagePdfDocumentReader pdfReader = new PagePdfDocumentReader(pdfResource);
        List<Document> documents = pdfReader.get();
        // 转换：按Token拆分
        TokenTextSplitter splitter = new TokenTextSplitter();
        List<Document> chunks = splitter.apply(documents);
        // 加载: 存储到ES向量数据库
        vectorStore.add(chunks);
    }

    @Test
    public void add4() {
        MarkdownDocumentReaderConfig config = MarkdownDocumentReaderConfig.builder()
                .withHorizontalRuleCreateDocument(true)
                .withIncludeCodeBlock(false)
                .withIncludeBlockquote(false)
                .withAdditionalMetadata("filename", markdownResource.getFilename())
                .build();
        MarkdownDocumentReader markdownDocumentReader = new
                MarkdownDocumentReader(markdownResource, config);
        List<Document> documents = markdownDocumentReader.get();
        // 转换：按Token拆分
        TokenTextSplitter splitter = new TokenTextSplitter();
        List<Document> chunks = splitter.apply(documents);
        // 加载: 存储到ES向量数据库
        vectorStore.add(chunks);
    }
}
